import os
import fitz  # PyMuPDF
import argparse
import numpy as np
import cv2
import tempfile
from rapidocr_onnxruntime import RapidOCR
from fix_torch import fix_torch_dll_path
from utils.logger import logger

fix_torch_dll_path()

class RapidOCRExtractor:
    def __init__(self, use_gpu=False, lang="ch"):
        """
        初始化RapidOCR提取器
        
        Args:
            use_gpu (bool): 是否使用GPU（RapidOCR会自动检测）
            lang (str): 语言，'ch'表示中文，'en'表示英文
        """
        # 初始化RapidOCR引擎
        self.ocr = RapidOCR()
        #logger.info(f"RapidOCR初始化完成，使用{'GPU' if use_gpu else 'CPU'}，语言：{lang}")
    
    def extract_from_pdf(self, page):
        """
        从PDF文件中提取字符和坐标信息
        
        Args:
            pdf_path (str): PDF文件路径
            output_dir (str, optional): 输出结果目录
            
        Returns:
            dict: 每页的OCR结果，格式为 {页码: {'chars': [{...}, {...}], 'text': str}}
        """
        
        results = {}
        
        # 获取页面信息
        page_width = page.rect.width
        page_height = page.rect.height
        
        # 获取页面图像，使用2x的分辨率提高OCR准确度
        zoom_factor = 2
        mat = fitz.Matrix(zoom_factor, zoom_factor)
        pix = page.get_pixmap(matrix=mat)
        
        # 创建临时文件保存图像
        temp_fd, temp_path = tempfile.mkstemp(suffix='.png')
        os.close(temp_fd)
        
        try:
            # 保存为图像文件
            pix.save(temp_path)
            
            # 使用RapidOCR进行识别
            ocr_result, elapse_time = self.ocr(temp_path, return_word_box=True)
            if isinstance(elapse_time, (list, tuple)):
                logger.info("OCR完成")
            else:
                logger.info(f"OCR完成，耗时: {float(elapse_time):.3f}秒")
            
            if ocr_result:
                # 处理OCR结果，提取单字符级别的坐标和文本
                page_text = ""
                page_chars = []
                
                for item in ocr_result:
                    # 检查OCR结果项的格式
                    if not isinstance(item, (list, tuple)):
                        logger.info(f"OCR结果项格式异常: {item}")
                        continue
                        
                    # 获取文本和置信度
                    text = item[1] if len(item) > 1 else ""
                    confidence = item[2] if len(item) > 2 else 0.0
                    
                    # 添加到页面文本
                    page_text += text + " "
                    
                    # 检查是否包含单字符级别的信息
                    if len(item) >= 5:  # 完整的结果应该有5个元素
                        char_boxes = item[3]  # 单字四点坐标列表
                        char_texts = item[4]  # 单字列表
                        
                        # 处理每个字符
                        for i, (box, char) in enumerate(zip(char_boxes, char_texts)):
                            try:
                                # 从四点坐标中提取矩形坐标
                                x_coords = [p[0] for p in box]
                                y_coords = [p[1] for p in box]
                                
                                # 原始坐标（以像素为单位，基于放大后的图像）
                                x0_pixel = float(min(x_coords))
                                y0_pixel = float(min(y_coords))
                                x1_pixel = float(max(x_coords))
                                y1_pixel = float(max(y_coords))
                                
                                # 将像素坐标转换为PDF坐标系统（考虑放大因子）
                                # PDF坐标是基于72dpi，需要除以放大因子进行转换
                                x0 = x0_pixel / zoom_factor
                                y0 = y0_pixel / zoom_factor
                                x1 = x1_pixel / zoom_factor
                                y1 = y1_pixel / zoom_factor
                                
                                # 添加到字符列表
                                page_chars.append({
                                    'bbox': [x0, y0, x1, y1],
                                    'char': char,
                                    'confidence': float(confidence)
                                })
                            except Exception as e:
                                logger.info(f"处理字符时出错: {e}")
                    else:
                        # 如果没有单字符级别的信息，打印警告
                        logger.info(f"警告: OCR结果缺少单字符信息: {item}")
                
                # 保存页面结果
                results = {
                    'chars': page_chars,
                    'text': page_text.strip()   
                }
                
                #logger.info(f"识别成功，共识别 {len(page_chars)} 个字符")
            else:
                #logger.info(f"识别失败，未找到文本")
                results = {'chars': [], 'text': ""}
            
        except Exception as e:
            #logger.info(f"处理时发生错误: {str(e)}")
            results = {'chars': [], 'text': ""}
        
        finally:
            # 删除临时文件
            if os.path.exists(temp_path):
                os.remove(temp_path)
        
        return results
    

def main():
    """主函数，处理命令行参数并调用OCR提取功能"""
    # 默认参数
    default_pdf_path = "./data/图像型测试.pdf"
    default_output_dir = "./output"
    
    parser = argparse.ArgumentParser(description="使用RapidOCR从PDF提取文本和坐标信息")
    parser.add_argument("--pdf", default=default_pdf_path, help=f"PDF文件路径，默认为{default_pdf_path}")
    parser.add_argument("--output", default=default_output_dir, help=f"输出结果目录，默认为{default_output_dir}")
    parser.add_argument("--use_gpu", action="store_true", help="是否使用GPU")
    parser.add_argument("--lang", default="ch", choices=["ch", "en"], help="识别语言，ch表示中文，en表示英文")
    
    args = parser.parse_args()
    
    # 创建OCR提取器
    extractor = RapidOCRExtractor(use_gpu=args.use_gpu, lang=args.lang)
    
    # 提取PDF中的文本和坐标信息
    results = extractor.extract_from_pdf(args.pdf)
    
    # 显示结果示例
    if results:
        for page_idx, page_data in results.items():
            if page_data['text']:
                logger.info(f"第 {page_idx + 1} 页文本: {page_data['text'][:100]}...")
            logger.info(f"第 {page_idx + 1} 页识别字符数: {len(page_data['chars'])}")

if __name__ == "__main__":
    main() 